//用于测试设备的IO性能，确定模型的性能参数
//Usage: io_profiler -t 1 -n 10000 -b 1 -s 4
//-t指线程数，-n指每个线程发起多少次IO，-b指一批最大发起多少个IO，-s指IO大小，单位KB
#include <iostream>
#include <thread> 
#include <random>
#include <sys/stat.h>
#include "../io_engine.h"
#include "../annlite_utils.h"
#include "statis.h"

using namespace std; 
 
//阻塞等待的IO，即发起一批请求后等待这一批请求全部完成后才会发起下一批
int io_profiler_sync(IOUringManager_ASYNC& io_manager, std::vector<void*>& thread_buf_list, long long total_io_count, long long per_io_size, int batch_size, long long range_start, long long range_end, Profiler_Statis& statis, int thread_id){ 
    // cout << "Thread " << thread_id << " started" << endl;
    //检查buf的size是否大于等于batch_size
    if(thread_buf_list.size() < (size_t)batch_size){
        cout << "Error: buf size:" << (int)thread_buf_list.size() << " is less than batch size:" << batch_size << endl;
        return -1;
    }
    // cout << "buf size:" << thread_buf_list.size() << ", batch size:" << batch_size << endl;
    //随机数生成器
    std::random_device rd;
    std::mt19937_64 gen(rd());
    std::uniform_int_distribution<long long> dis(range_start, range_end-per_io_size);

    int undispatch_io = total_io_count; 
    vector<long long> aligned_offset_list;
    vector<long long> aligned_length_list;
    vector<int> ids;  

    //记录开始时间
    auto start_time = statis.get_time();

    int batch_round = 0;
    while(1){ 
        if(undispatch_io <= 0){
            break;
        }
        //计算当前批次能够发起的IO数量
        int batch_io_count = min(batch_size, undispatch_io);
        //构造请求
        for(int i = 0; i < batch_io_count; i++){
            long long offset = dis(gen);
            offset = aligned_offset_in_file(offset); //获取起始地址
            aligned_offset_list.push_back(offset);
            aligned_length_list.push_back(per_io_size);
            ids.push_back(i);
        }
        // print_array(aligned_offset_list, aligned_offset_list.size());
        // print_array(aligned_length_list, aligned_length_list.size());
        // print_array(ids, ids.size());

        //提交请求
        if (!io_manager.submitRequests(thread_buf_list, aligned_offset_list, aligned_length_list, ids, batch_io_count)) {
            std::cout << "Failed to submit IO requests." << std::endl;
            exit(0);
        }
        //等待IO完成
        int complete_request = 0;
        while(1){
            std::vector<int> completedIndices;
            int complete_index_count = io_manager.getCompletedResultIds(completedIndices, batch_io_count);
            for(int i=0; i<complete_index_count; i++){
                int complete_index = completedIndices[i];
                // cout << "Thread " << thread_id << " completed IO " << complete_index << endl;
                char* data = (char*)(thread_buf_list[complete_index]);
                data = data + 1;
                // cout << "data[0]:" << ((float*)data)[0] << endl;
            }
            complete_request += complete_index_count;
            if (complete_request >= batch_io_count) {
                // std::cout << "Thread " << thread_id << " batch " << batch_round << "'s "<< complete_index_count <<" requests completed." << std::endl;
                break;
            }
        }
        undispatch_io -= complete_request;
        batch_round++;
        //统计IO次数
        statis.add_io_count(thread_id, complete_request);
        //统计IO大小
        statis.add_io_size(thread_id, complete_request*per_io_size);
    }

    //记录结束时间
    auto end_time = statis.get_time();
    double time = statis.elapsed_time(start_time, end_time);
    statis.add_time(thread_id, time);

    return 0;
}

int main(int argc, char* argv[])
{
    // 默认参数值
    string test_file_path = "/home/ljl/Code/dataset/vector-ssd/sift1m/sift_1m_base.fbin";
    int thread_count = 1;
    long long total_io_count = 1000;
    long long per_io_size = 4*KB;
    int batch_size = 2;
    int io_mode = 0;

    // 解析命令行参数
    for (int i = 1; i < argc; ++i) {
        string arg = argv[i];
        if (arg == "-f" && i + 1 < argc) {
            test_file_path = argv[++i];
        } else if (arg == "-t" && i + 1 < argc) {
            thread_count = stoi(argv[++i]);
        } else if (arg == "-n" && i + 1 < argc) {
            total_io_count = stoll(argv[++i]);
        } else if (arg == "-s" && i + 1 < argc) {
            per_io_size = stoll(argv[++i]) * KB;
        } else if (arg == "-b" && i + 1 < argc) {
            batch_size = stoi(argv[++i]);
        } else if (arg == "-m" && i + 1 < argc) {
            io_mode = stoi(argv[++i]);
        } else if (arg == "-h") {
            cout << "Usage: " << argv[0] << " [-f file_path] [-t thread_count] [-n total_io_count] "
                 << "[-s per_io_size_MB] [-b batch_size] [-m io_mode]" << endl;
            return 0;
        }
    }

    cout << "BLOCK_SIZE: " << ANN_BLOCK_SIZE << endl;
    cout << "Total IO count: " << total_io_count << endl;
    cout << "Per IO size: " << per_io_size << endl;
    cout << "Batch size: " << batch_size << endl;
    cout << "Thread num: " << thread_count << endl;
    //判断test_file_path是否存在，并打印文件的大小
    struct stat file_stat;
    if (stat(test_file_path.c_str(), &file_stat) != 0) {
        std::cerr << "Failed to get file size: " << test_file_path << std::endl;
        exit(1);
    }
    long long file_size = file_stat.st_size;
    cout << "File path: " << test_file_path << endl;
    cout << "File size: " << file_size << endl;

    std::vector<IOUringManager_ASYNC> io_manager_array;
    io_manager_array.reserve(thread_count);
    std::vector<std::vector<void*>> io_manager_buf_array_2d;
    std::vector<int> fd_array;
    Profiler_Statis statis(thread_count);

    for (int i = 0; i < thread_count; i++) {
        //初始化文件句柄
        int fd = open(test_file_path.c_str(), O_RDONLY | O_DIRECT);
        if (fd == -1) {
            std::cerr << "Failed to open file: " << test_file_path << std::endl;
            exit(1);
        }
        fd_array.push_back(fd);
        //初始化io_manager
        io_manager_array.emplace_back(fd_array[i], batch_size);
        //为每个线程分配buffer
        std::vector<void*> io_manager_buf_array_1d;
        for(int j=0; j < batch_size; j++){ //每个IO的buffer
            io_manager_buf_array_1d.push_back(get_aligned_buffer(per_io_size));
        }
        io_manager_buf_array_2d.push_back(io_manager_buf_array_1d);
    }

    //开始测试
    cout << "Start IO performance profile" << endl;
    statis.start(); 

    std::vector<std::thread> threads;
    for (int i = 0; i < thread_count; i++) {
        if(io_mode == 0){
            // 修改捕获方式，按值捕获需要的变量
            threads.emplace_back([=, &io_manager_array, &io_manager_buf_array_2d, &statis]() {
                io_profiler_sync(io_manager_array[i], io_manager_buf_array_2d[i], total_io_count, per_io_size, batch_size, 0, file_size, statis, i);
            });
        }
    }
    
    for (auto& t : threads) {
        t.join();
    }

    statis.end();
    //输出性能数据
    statis.io_print();

    for(int i = 0; i < thread_count; i++){
        if(fd_array[i]!=-1)
            close(fd_array[i]);
        for(size_t j=0; j<io_manager_buf_array_2d[i].size(); j++){
            free(io_manager_buf_array_2d[i][j]);
        }
    }

    return 0;
}